#include <string.h>
#include "cover.h"

// Global cover
Cover * G_COVER = NULL;

void IndexRollingHash(const string & s, vector<uint32> * ret) {
  ret->clear();
  uint32 last_val = 0;
  for (uint32 i=0; i<s.size(); i++) {
    last_val = last_val * kRollingHashK + s[i];
    ret->push_back(last_val);
  }
}

// compute k^n mod 2^32
uint32 UInt32Pow(uint32 k, uint32 n) {
  uint32 ret = 1;
  while (n) {
    if (n & 1) ret *= k;
    n >>= 1;
    k *= k;
  }
  return ret;
}

uint32 ComputeRollingHash(const vector<uint32> &a, int begin, int end) {
  return a[end-1] - ((begin==0) ? 0 : a[begin-1]) 
    * UInt32Pow(kRollingHashK, end - begin);
}


SubstringIndex::SubstringIndex(const string &s) {
    hash_position_.resize(kSearchHashBuckets);
    s_ = s;
    IndexRollingHash(s, &rolling_hash_);
    for (uint32 window_size = 1; window_size < s.size(); window_size *= 2) {
      for (int start = 0; start + window_size <= s.size(); start +=window_size){
	int end = start + window_size;
	uint32 h = ComputeRollingHash(rolling_hash_, start, end);
	hash_position_[h % kSearchHashBuckets].push_back(start);
      }
    }
  }
  
void SubstringIndex::Search(const string & s, vector<uint32> * ret) {
  ret->clear();
    vector<uint32> rh;
    IndexRollingHash(s, &rh);
    uint32 s_hash = ComputeRollingHash(rh, 0, s.size());
    uint32 window_size = 1; 
    while (window_size * 4  <= s.size() + 1 ) window_size *= 2;
    for (uint32 start = 0; start < window_size; start++) {
      uint32 h = ComputeRollingHash(rh, start, start+window_size);
      const vector<int>& positions = hash_position_[h % kSearchHashBuckets];
      for (uint32 i=0; i<positions.size(); i++) {
	int pos = positions[i];
	if (ComputeRollingHash(rolling_hash_, pos-start, pos-start + s.size())
	    == s_hash)
	  if (strncmp(s.c_str(), s_.c_str() + pos-start, s.size()) == 0)
	    ret->push_back(pos-start);
      }
    }
}

string Element::ToString() {
  ostringstream o;
  o << "address:" << this << " " 
    << s_ << " solution_:" << solution_ << " use_count_:" << use_count_;
  return o.str();
}

Element::Element(const Element& e) {
  s_ = e.s_;
  solution_ = e.solution_;
  use_count_ = e.use_count_;
}

Element::Element(Solution * sol, string s) {
  solution_ = sol; 
  s_ = s;
  use_count_ = 0;
}
string Solution::ToString() {
  ostringstream o;
  o << "address:" << this;
  o << " use_count_:" << use_count_;
  o << " prior_ln_likelihood_:" << prior_ln_likelihood_ << endl;
  for (uint32 i=0; i<elements_.size(); i++) {
    o << elements_[i].ToString(); o << endl; }
  return o.str();
}

Solution::Solution() {
  use_count_ = 0;
  prior_ln_likelihood_ = 0;
}

void Solution::AddElement(string s) {
  elements_.push_back(Element(this, s));
  // cout << elements_.back().solution_ << endl;
}

void Cover::AddElementAt(uint32 pos, Element * e, bool add_change) { 
  e->use_count_++;
  sum_ln_factorial_counts_ += log(e->use_count_);
  
  // We keep track of the number of elements possible to reference
  // using our current solutions that are used
  if (e->solution_->use_count_ == 0) {
    solutions_.insert(e->solution_);
    total_num_elements_ += e->solution_->elements_.size();
    solution_ln_likelihood_ += e->solution_->prior_ln_likelihood_;
  }
  e->solution_->use_count_++;
  elements_.insert(make_pair(pos, e));
  if (add_change) {
    added_.push_back(make_pair(pos, e));
  }
  
}

void Cover::RemoveElementAt(uint32 pos, Element * e, bool add_change) {
  sum_ln_factorial_counts_ -= log(e->use_count_);
  e->use_count_--;
  e->solution_->use_count_--;
  if (e->solution_->use_count_ == 0) {
    solutions_.erase(e->solution_);
    if (add_change) solutions_removed_.push_back(e->solution_);
    total_num_elements_ -= e->solution_->elements_.size();
    solution_ln_likelihood_ -= e->solution_->prior_ln_likelihood_;
  }
  elements_.erase(pos);
  if (add_change) {
    removed_.push_back(make_pair(pos, e));
  }
}

Cover::Cover(SubstringIndex * si) {
  sum_ln_factorial_counts_ = 0;
  si_ = si;
  solution_ln_likelihood_ = 0;
  total_num_elements_ = 0;
  
  // First create an initial solution with only the elements used in that string
  alphabet_ = new Solution();
  int element_map[256];
  memset(element_map, 0, 256 * sizeof(int));
  string& s = si_->s_;
  for (uint32 i=0; i<s.size(); i++) {
    uint32 c = uint32((unsigned char)s[i]);
    if (element_map[c] == 0) {
      // cout << "c:" << c << endl;
      // cout << "Adding element for " << s.substr(i, 1) << endl;
      element_map[c] = alphabet_->elements_.size();
      alphabet_->AddElement(s.substr(i, 1));
    }
  }
  
  // cout << "SOLUTION!\n";
  // cout << alphabet_->ToString() << endl;
  
  // Next cover the string with that solution 
  for (uint32 i=0; i<s.size(); i++) {
    uint32 c = uint32((unsigned char)s[i]);
    // cout << "covering c:" << c << endl;
    AddElementAt(i, &alphabet_->elements_[element_map[c]], false);
  }
}

// Try adding a solution, report back the improvement, difference?
double Cover::TryAddSolution(Solution * s,
			     vector<Solution *> * solutions_removed) {
  double old_score = Score();
  added_.clear();
  removed_.clear();
  solutions_removed_.clear();
  for (uint32 c=0; c<s->elements_.size(); c++) {
    TryAddElement(&(s->elements_[c]));
  }
  if (Score() > old_score) {
    cout << "improvement! old_score = " << old_score << " new score = " 
	 << Score() << endl;
    cout << "diff = " << Score() - old_score << " nats" << endl;
    cout << s->ToString() << endl;
    return Score() - old_score;
  } else {
    for (uint32 i=0; i<added_.size(); i++) 
      RemoveElementAt(added_[i].first, added_[i].second, false);
    for (uint32 i=0; i<removed_.size(); i++) 
      AddElementAt(removed_[i].first, removed_[i].second, false);
    return 0.0;
  }
}

void Cover::TryAddElement(Element * e) {
  // First search for this element everywhere
  vector<uint32> hits;
  si_->Search(e->s_, &hits);
  for (uint32 c=0; c<hits.size(); c++)
    TryAddSegment(e, hits[c]);
}

// The segment needs to replace at least 1 other segment to be viable
void Cover::TryAddSegment(Element * e, uint32 pos) {
  int segments_removed = 0;
  uint32 endpos = pos + e->s_.size();
  map<uint32, Element*>::iterator 
    find = elements_.lower_bound(pos);
  while ((find != elements_.end()) &&
	 ((find->first + find->second->s_.size()) <= endpos)) {
    segments_removed++;
    uint32 remove_pos = find->first;
    Element * remove_element = find->second;
    find++;
    RemoveElementAt(remove_pos, remove_element, true);
  }
  
  // If any segments were removed add this segment
  if (segments_removed > 0) {
    AddElementAt(pos, e, true);
  }
}

void InitializeGlobalCover() {
  string corpus; 
  ReadFileIntoString(&corpus, "wiki.1m");
  cout << "read string\n";

  SubstringIndex * index = new SubstringIndex(corpus);
  cout << "formed index\n";

  G_COVER = new Cover(index);
  cout << "Created global cover initial score:" << G_COVER->Score() << endl;
}

int StringMatchMain() {
  
  InitializeGlobalCover();

  /*
  string s;
  while (cin >> s) {
    vector<uint32> positions;
    G_COVER.si_->Search(s, &positions);
    cout << "matches: ";
    for (uint32 i=0; i<positions.size(); i++) cout << positions[i] << " ";
    cout << endl;
  }
  */

  string s;
  while (cin >> s) {
    Solution * sol = new Solution();
    sol->AddElement(s);
    vector<Solution *> removed;
    G_COVER->TryAddSolution(sol, &removed);
    cout << "new cover score:" << G_COVER->Score() << endl;
  }

  return 0;
}
